from __future__ import absolute_import

from __future__ import division
from __future__ import print_function

import sys
import os
import argparse
import numpy as np
import mxnet as mx
import cv2
from . import face_align
from PIL import Image
from faceapp.models import face
from faceapp.serializers import FaceSerializer
import base64
from rsa_face.RSA import RSA


def do_flip(data):
    for idx in range(data.shape[0]):
        data[idx, :, :] = np.fliplr(data[idx, :, :])


def get_model(ctx, image_size, prefix, epoch, layer):
    print('loading', prefix, epoch)
    sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
    all_layers = sym.get_internals()
    sym = all_layers[layer + '_output']
    model = mx.mod.Module(symbol=sym, context=ctx, label_names=None)
    model.bind(data_shapes=[('data', (1, 3, image_size[0], image_size[1]))])
    model.set_params(arg_params, aux_params)
    return model


class FaceModel:
    def __init__(self, ctx_id, model_prefix, model_epoch):
        self._feature_len = 256
        self.detector = RSA(max_img=512)

        if ctx_id >= 0:
            ctx = mx.gpu(ctx_id)
        else:
            ctx = mx.cpu()
        image_size = (112, 112)
        self.model = get_model(
            ctx, image_size, model_prefix, model_epoch, 'fc1')
        self.image_size = image_size
        queryset = face.objects.all()
        serializer = FaceSerializer(queryset, many=True)
        self.instances = serializer.instance
        self.update_db()


    def file2img(self, file):
        rgbImg = Image.open(file)
        cvImg = np.array(rgbImg)
        img = cv2.cvtColor(cvImg, cv2.COLOR_RGB2BGR)
        return img


    def get_input(self, face_img):
        bboxes, pts5s = self.detector.detect(face_img)
        if bboxes is None:
            return None
        if bboxes.shape[0]==0:
            return None
        bbox = bboxes[0, 0:4]
        pts5 = pts5s[0, 0:10].reshape(5, 2)
        nimg = face_align.norm_crop(face_img, pts5)
        return nimg

    def get_feature(self, face_img, align=False, base64img=False):
        if base64img:
            try:  
                newbyte = base64.b64decode(face_img.encode('utf-8'))
                newarray = np.frombuffer(newbyte, np.uint8)
                face_img = cv2.imdecode(newarray, cv2.IMREAD_COLOR)
                cv2.imwrite('face_img.jpg', face_img)
            except:
                print('base64 decode error')
                return None

        if not align:
            face_img = self.get_input(face_img)
            if face_img is None:
                return None
        a = cv2.cvtColor(face_img, cv2.COLOR_BGR2RGB)
        a = np.transpose(a, (2, 0, 1))
        input_blob = np.expand_dims(a, axis=0)
        data = mx.nd.array(input_blob)
        db = mx.io.DataBatch(data=(data,))
        self.model.forward(db, is_train=False)
        emb = self.model.get_outputs()[0].asnumpy()[0]
        norm = np.sqrt(np.sum(emb * emb) + 0.00001)
        emb /= norm
        return emb

    def compare_1v1(self, face_img1, face_img2):
        fea1 = self.get_feature(face_img1)
        fea2 = self.get_feature(face_img2)
        if fea1 is None or fea2 is None:
            return None
        sim = np.dot(fea1, fea2) / \
            (np.linalg.norm(fea1) * np.linalg.norm(fea2))
        return (sim+1)/2.0

    def compare_1vn(self, fea):
        if self._face_db_fea is None:
            return None
        if fea is None:
            return None
        fea = fea.reshape(self._feature_len, 1)
        res = np.matmul(self._face_db_fea, fea)
        res = res[:, 0]
        max_value = np.max(res)
        max_value = (max_value+1)/2.0
        max_ind = int(np.argmax(res))
        res = {"score":max_value, "instance": self.instances[max_ind]}
        return res

    @property
    def face_db_fea(self):
        return self._face_db_fea

    @property
    def face_db_name(self):
        return self._face_db_name

    @property
    def feature_len(self):
        return self._feature_len

    def update_db(self):
        queryset = face.objects.all()
        serializer = FaceSerializer(queryset, many=True)
        self.instances = serializer.instance
        db_fea = []
        db_name = []
        for (ind, instance) in enumerate(self.instances):
            if len(instance.feature)%4!=0:
                continue
            fea = base64.b64decode(instance.feature.encode('utf-8'))

            if len(fea)==self._feature_len*4:
                fea = np.frombuffer(fea, np.float32).reshape(1, self._feature_len)
            else:
                continue
            db_name.append(instance.name)
            db_fea.append(fea)
        if len(db_fea) == 0:
            self._face_db_fea = None
        else:
            self._face_db_fea = np.concatenate(db_fea, 0)
        self._face_db_name = db_name
